import datetime
import os
import time
from pathlib import Path
import argparse
from omegaconf import OmegaConf

from mindspore import nn
from mindspore.train.model import Model
from mindspore import load_param_into_net, load_checkpoint

from models import build_model
import misc_utils.misc as misc
from misc_utils.logger import create_logger
from data_utils import build_eval_dataset


def main(args):
    args.rank, args.world_size = misc.init_context(args)
    logger = create_logger(args)
    logger.info(f'job dir: {args.output_dir}')

    model = build_model(args)
    logger.info(f'Loading checkpoint: {args.ckpt}')
    ckpt = load_checkpoint(args.ckpt)
    param_not_load, _ = load_param_into_net(model, ckpt)
    logger.info(f'Param not loaded: {param_not_load}')
    num_params = sum([param.size for param in model.get_parameters() if param.requires_grad])
    logger.info(f'Number of parameters: {num_params}')
    model.set_train(False)

    dataset_collection = build_eval_dataset(args)
    eval_metrics = {
        'loss': nn.metrics.Loss()
    }
    trainer = Model(model, eval_network=model, metrics=eval_metrics)

    for dset_name, dataset in dataset_collection.items():
        per_step_size = dataset.get_dataset_size()
        logger.info(f'Eval {dset_name} with {per_step_size} steps')

        start_time = time.time()
        result = trainer.eval(valid_dataset=dataset, dataset_sink_mode=False)
        total_time = time.time() - start_time
        total_time_str = str(datetime.timedelta(seconds=int(total_time)))
        logger.info(result)
        logger.info(f'Time: {total_time_str}\n')


def create_parser():
    parser = argparse.ArgumentParser('Med Foundation Models')
    # parser.add_argument("--config", default="", metavar="FILE", help="path to config file")
    parser.add_argument('--ckpt', type=str, default='')
    parser.add_argument('--output_dir', default='/home/yueyang/output')
    parser.add_argument('--exp_name', type=str, default='debug')
    parser.add_argument(
        "opts",
        default=None,
        nargs=argparse.REMAINDER,
    )
    return parser


if __name__ == '__main__':
    args = create_parser().parse_args()
    args.output_dir = os.path.join(args.output_dir, 'eval_' + args.exp_name)
    print('Output Dir:', args.output_dir)
    Path(args.output_dir).mkdir(parents=True, exist_ok=True)

    args.opts += [f"output_dir={args.output_dir}", f'ckpt={args.ckpt}']
    # default_cfg = OmegaConf.load('configs/default.yaml')
    cfg = OmegaConf.load(os.path.join(os.path.dirname(args.ckpt), 'config.yaml'))
    cfg = OmegaConf.merge(cfg, OmegaConf.from_cli(args.opts))
    with open(os.path.join(args.output_dir, 'config.yaml'), "w") as f:
        OmegaConf.save(config=cfg, f=f)
    main(cfg)
